import argparse
import numpy as np
import torch

parser = argparse.ArgumentParser(description='Chinese Text Classification')
parser.add_argument('--model', type=str, required=True, help='choose a model: TextCNN, TextRNN, FastText, TextRCNN, TextRNN_Att, DPCNN, Transformer')
parser.add_argument('--embedding', default='random', type=str, help='random or embedding_SougouNews.npz or else')
parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')
parser.add_argument('--dataset', default='THUCNews', type=str, help='数据文件夹')
parser.add_argument('--seed', default=123, type=int, help='')

# 共用
parser.add_argument('--dropout', default=0.2, type=float, help='')
parser.add_argument('--require_improvement', default=1000, type=int, help='若超过1000batch效果还没提升，则提前结束训练')
parser.add_argument('--num_epochs', default=20, type=int, help='')
parser.add_argument('--batch_size', default=128, type=int, help='')
parser.add_argument('--pad_size', default=32, type=int, help='')
parser.add_argument('--learning_rate', default=1e-3, type=float, help='Transformer默认5e-4')

# DPCNN/TextCNN
parser.add_argument('--num_filters', default=256, type=int, help='卷积核数量.DPCNN/TextCNN')

# FastText
parser.add_argument('--hidden_size', default=256, type=int, help='FastText/TextRCNN/TextRNN')
parser.add_argument('--n_gram_vocab', default=250499, type=int, help='ngram 词表大小/FastText')

# TextRCNN
parser.add_argument('--num_layers', default=1, type=int, help='TextRCNN/TextRNN_Att是2')

# Transformer
parser.add_argument('--dim_model', default=300, type=int, help='')
parser.add_argument('--hidden', default=1024, type=int, help='')
parser.add_argument('--last_hidden', default=512, type=int, help='')
parser.add_argument('--num_head', default=5, type=int, help='')
parser.add_argument('--num_encoder', default=2, type=int, help='')

# TextRNN_Att
parser.add_argument('--hidden_size2', default=64, type=int, help='TextRNN_Att')

# bert
parser.add_argument('--bert_path', default="./pretrain_models/bert-base-chinese", type=str, help='')
parser.add_argument('--freeze_to_layer', default=None, type=str, help='例如 bert.pooler.dense.bias')

# 多标签
parser.add_argument('--multi_labels', default=False, type=bool, help='重要!!是否是多标签多分类')
parser.add_argument('--print_per_batch', default=100, type=int, help='多少batch打印一次')
parser.add_argument('--threshold', default=0.5, type=float, help='阈值')
parser.add_argument('--MLB_path', default="mlb.pkl", type=str, help='')

# 推理
parser.add_argument('--infer_mode', default=False, type=bool, help='')
parser.add_argument('--infer_combine', default=False, type=bool, help='转换标签并合并到同一个文件')

# Sanic
parser.add_argument('--port', default=6060, type=int, help='')
parser.add_argument('--ip', default="0.0.0.0", type=str, help='')


args = parser.parse_args()

np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
torch.backends.cudnn.deterministic = True  # 保证每次结果一样

assert args.model in ["TextCNN", "TextRNN", "FastText", "TextRCNN", "TextRNN_Att", "DPCNN", "Transformer",
                        "bert", "bert_CNN", "bert_DPCNN", "bert_RCNN", "bert_RNN"]